﻿
#------------------------------
# Imports
import bpy

from myxml import *
    
#---------------------------------------------   
    

# Obj should be an armature object
def write(obj, filename):
    
    if(obj.type != 'MESH'):
        print('Mesh: selected object not a mesh object')
        return

    file = open(filename, 'w')

    def w(text) :
        file.write(text)
    
    write_object(obj, w)
    
    file.close()
    
    print("Object '" + obj.name + "' : mesh exported")
    
def write_object(obj, w) :     

    w(otag('mesh'))
    
    # Obtain mesh to be exported (based on export_x.py)
    
    # We disable the armature modifiers, to have access to the untransformed
    # vertices (rest pose)
    deactivated_modifiers = [m for m in obj.modifiers if m.type == 'ARMATURE' and m.show_render]
    
    for m in deactivated_modifiers:
        m.show_render = False
            
    # We bake the modifier stack of the mesh object into new mesh data.
    # This baked data is exported.
    # see function to_mesh() in 
    #     http://www.blender.org/documentation/blender_python_api_2_67_1/bpy.types.Object.html#bpy.types.Object
    mesh = obj.to_mesh(bpy.context.scene, True, 'RENDER')
    
    # TODO: provide mechanism to specify which uv map should be used
    uv_layer = mesh.uv_layers.active
        
    (uv_to_global,refers,to_orig,orig_to_global) = create_vertex_clones(mesh, uv_layer)
    
    write_verts(mesh, to_orig, refers, w)
    write_submeshes(mesh, uv_to_global, uv_layer, w)
    # write_polys(mesh, w)    
    #write_uvmaps(mesh, w)
    
   
    # Pass dictionary with (group id -> group name) mappings.
    groups = {}
    for gr in obj.vertex_groups :
        groups[gr.index] = gr.name
        
    # write_vert_groups(mesh, groups, w)
    write_skinweights(mesh, groups, deactivated_modifiers, orig_to_global, w)
    
    # Delete the baked mesh
    bpy.data.meshes.remove(mesh)
    
    # enable previously disabled modifiers
    for m in deactivated_modifiers:
        m.show_render = True

    w(ctag('mesh'))

#---------------------------------------------

# Function global_to_orig maps new vertex ids to original
# vertex ids
def write_verts(data, global_to_orig, refers, w) :
    
    w(otag('vertices', [('count', len(global_to_orig))]))
    
    for i in range(0, len(global_to_orig)) :
        orig_id,uv = global_to_orig[i]
        
        v = data.vertices[orig_id]
    
        # Write identifier
        if(i == refers[i]) : # refers to itself 
            w(otag('vert', [('id', i)]))           
        else:
            w(otag('vert', [('id', i), \
                            ('ref', refers[i])]))

        w(octag('coords', zip('xyz', v.co)))
        w(octag('normal', zip('xyz', v.normal)))
        
        w(octag('tex', zip('uv', uv)))
        
        w(ctag('vert'))
        
    w(ctag('vertices'))
    
#---------------------------------------------

# Submeshes are based on the material identifier of the polygons.
def write_submeshes(mesh, to_global, uv_layer, w) :
    
    # First we collect all used material identifiers.
    # We map each material id to the number of polygons
    # labelled by it.
    material_ids = dict()
    
    for p in mesh.polygons:
        if material_ids.get(p.material_index) is None:
            material_ids[p.material_index] = 0
        material_ids[p.material_index] = material_ids[p.material_index] + 1
        
    w(otag('submeshes', [('count', len(material_ids))]))
        
    # Now iterate over the used materials, listing the polygons
    # for each submesh.    
    for mi in material_ids:
        w(otag('submesh', [('name', bpy.data.materials[mi].name), \
                           ('polycount', material_ids[mi])]))
        
        for p in mesh.polygons:
            if p.material_index == mi:
                w(otag('poly', [('id', p.index)]))     
        
                # Write vertices (these are in counter-clockwise order).
                # Using the uv_layer, we know which global index to 
                # use (i.e., which clone)
                for loop_index in p.loop_indices :                
                    vert_id = mesh.loops[loop_index].vertex_index
                    uv = uv_layer.data[loop_index].uv
                    
                    uvtup = (uv[0], uv[1])
                    
                    global_index = to_global[(vert_id, uvtup)]                                
                    
                    w(octag('vert', [('id', global_index)]))     
        
                w(ctag('poly'))
        
        w(ctag('submesh'))
        
    w(ctag('submeshes'))
    
    
#---------------------------------------------

def create_vertex_clones(mesh, uv_layer) :
    vertex_uvs = dict()
    
    # Give each vertex an empty set of uv coordinate pairs
    for v in mesh.vertices:
        vertex_uvs[v.index] = set()        
        
    # Based on the manual entry
    # http://www.blender.org/documentation/blender_python_api_2_65_5/bpy.types.Mesh.html#bpy.types.Mesh.tessface_uv_textures
    # The polygons reference loops.
    
    # Citation:
    #   "Each polygon reference a slice in the loop array, this way, 
    #   polygons do not store vertices or corner data such as UV’s directly, 
    #   only a reference to loops that the polygon uses."
    #
    # This corresponds well to the half-edge data structure. Ok, so each polygon
    # is a counter-clockwise sequence of vertices, connected by the loop. Each 
    # loop is basically a half-edge, that points to the next half-edge (I guess).
        
    # Now go over the polygons to give each vertex its
    # uv pairs. Along seam edges, vertices are typically given
    # different uv pairs by each face sharing the seam edge.
    for p in mesh.polygons:
        for loop_index in p.loop_indices :                
            vert_id = mesh.loops[loop_index].vertex_index
            
            uv = uv_layer.data[loop_index].uv
            
            vertex_uvs[vert_id].add((uv[0],uv[1]))

    
    # Now we duplicate vertices with more than one uv-pair
    to_global_ids = dict() # maps (vertex_id, uv-pair) to new global vertex index
    refers = [] # global ids to the global id that they clone
    global_to_orig = [] # global ids to original ids with uv-data (list for efficiency)
    orig_to_global = [] # original ids to all their clone ids (each entry is a list)
    
    # init orig_to_global
    for v in mesh.vertices:
        orig_to_global.append([])
    
    global_index = 0
    for vert_id, uv_set in vertex_uvs.items() :
        # Remember the first global index for this
        # vertex
        first = global_index
        
        # Put an arbitrary order on the uv-set
        for uv in uv_set:
            to_global_ids[(vert_id, uv)] = global_index
            
            refers.append(first)
            global_to_orig.append((vert_id,uv))
            orig_to_global[vert_id].append(global_index)
            
            global_index += 1
            
#     print("----------------")
#     print("----------------")
#     print(refers)
#     print("----------------")
#     print(to_global_ids)
    return (to_global_ids, refers, global_to_orig, orig_to_global)

#---------------------------------------------

def write_vert_groups(data, groups, w) :
        
    # The data is stored inversely:
    # each vertex has a list of group identifiers
    # So, for each group we first collect the
    # vertex identifiers. 
    group_to_vertices = {}
    for groupid in groups : 
        group_to_vertices[groupid] = list()  # Initialize to empty list
    
    num_verts = len(data.vertices)
    
    for i in range(0, num_verts) : 
        vertex = data.vertices[i]
        
        for group in vertex.groups :  # type(group) = bpy.types.VertexGroupElement
            group_id = group.group             
            
            # A bug probably: sometimes vertices contain nonexistent groups.
            # Hence the following test.
            if(group_to_vertices.get(group_id) == None) :
                pass
            else :
                group_to_vertices[group_id].append(i)            
    
    # Now write each group to file
    w(otag('groups'))
    
    for group_id in groups : 
        
        w(otag('group', [('id', group_id), \
                         ('name', groups[group_id]), \
                         ('size', len(group_to_vertices[group_id])) \
                         ]))
        
        # Actual vertices        
        for vert_id in group_to_vertices[group_id] :
            w(octag('vert', [('id', vert_id)]))
        
        w(ctag('group'))
    
    w(ctag('groups'))    
    
#---------------------------------------------
    
# groups maps group ids to group names
# We also pass the armature modifiers, so that multiple skins can be exported,
# one for each modifier.
def write_skinweights(mesh, groups, modifiers, orig_to_global, w):
    
    #print(orig_to_global)
    
    w(otag('skins', [('count', len(modifiers))]))
    
    for m in modifiers:
        
        w(otag('skin', [('name', m.name), \
                        ('normalized', 'false')]))
        
        # We only export bone weights for bones in the armature
        # object of this modifier
        bone_names = set([bone.name for bone in m.object.data.bones])
        
        # We map each bone to a list of weighted vertices for that bone
        influenced_verts = dict()
        for bone_name in bone_names:
            influenced_verts[bone_name] = []
    
        # Iterate over the vertices (in order). This ensures
        # that the list of vertices influenced by each bone
        # will be nicely ordered as well. 
        for v in mesh.vertices:            
            
            # Iterate over VertexGroupElements                    
            for vge in v.groups:                
                group_name = groups[vge.group]
                
                if group_name in bone_names:
                    # Because this                     
                    #w(octag('bone', [('name', group_name), ('weight', vge.weight)]))
                    
                    list = influenced_verts[group_name]
                    
                    # Put all clones of the vertex as being influenced by the bone
                    for clone_id in orig_to_global[v.index]:
                        list.append((clone_id,vge.weight))                    
            
        # Now we write the influenced clone vertices for each bone
        
        for bone_name, list in influenced_verts.items():
            if len(list) == 0:
                continue # do not write bones without influence
            
            w(otag('bone', [('name', bone_name)]))
            
            for clone_id, weight in list:
                w(octag('vert', [('id', clone_id), \
                                 ('weight', weight)]))            
            
            w(ctag('bone'))            
            
        w(ctag('skin'))
    
    w(ctag('skins'))


